#ifndef ELEMENT_H_
#define ELEMENT_H_

#include "Geometry.h"
#include <eigen3/Eigen/Dense>
#include <fstream>
#include <string>
#include <functional>
#include <iostream>
#include <iomanip>
using std::setw;
using std::setiosflags;

/**
 * @brief 自由度类，继承于Vertex，增加了编号这一成员
 * 
 * @tparam DIM 
 */
template<int DIM>
class Dofs;

using std::vector;
using std::string;
using std::ifstream;
using std::ofstream;
using std::function;
template<int DIM>
class Dofs : public Vertex<DIM>
{
private:
    size_t index;
public:
    Dofs() : Vertex<DIM>() { index = 0;};
    Dofs(vector<double> &vec) : Vertex<DIM>() { this->set_coordinate(vec), index = 0;};
    void set_coordinate(std::vector<double> _coord) { this->Vertex<DIM>::set_coordinate(_coord);};
    vector<double> get_coordinate() { return this->Vertex<DIM>::get_coordiante();};
    void set_index(size_t _index) { index = _index ;};
    size_t get_index() { return index; };
};

/**
 * @brief 单元类的基类，抽象类
 * 
 * @tparam DIM 
 */
template<int DIM>
class Element;

template<int DIM>
class Element
{
protected:
    size_t index;
    vector<Dofs<DIM> > dofs;
public:
    /**
     * @brief 设置单元编号
     * 
     * @param _index 
     */
    void set_index(size_t _index) { index = _index; };
    size_t get_index() { return index; };
    /**
     * @brief 设置第_i个自由度的信息
     * 
     * @param _d 
     * @param _i 
     */
    virtual void set_dofs(Dofs<DIM> &_d, int _i) { dofs[_i] = _d; };
    virtual void set_dofs(vector<double> &vec, int _i, int _index)
    {
        dofs[_i].set_coordinate(vec);
        dofs[_i].set_index(_index);
    }
    virtual vector<vector<double> > get_vertex_coordinate() = 0;
    virtual const vector<Dofs<DIM> >& get_dofs() const { return dofs; };
    /**
     * @brief 第i和j个基函数对刚度矩阵的贡献，贡献位置为get_i_dofs(_i).get_index(), get_i_dofs(_j).get_index()
     * 
     * @param weight_points 输入的积分点
     * @param quad_weights 输入的积分权重
     * @param _i 
     * @param _j 
     * @return double 
     */
     Dofs<DIM> & get_i_dofs (size_t _i)   { return dofs[_i]; };
    virtual double con_Matrix(vector<Vertex<DIM> > &weight_points, vector<double> &quad_weights, int _i, int _j) = 0 ;
};

//关于三角形单元，待完成
/////////////////////////////////////////////////////////////////////////////////////////////////////////
class P_Element;

class P_Element : public Element<2>, public Triangle
{
public:
    P_Element() :  Triangle() { this->index = -1;};
    vector<vector<double> > get_vertex_coordinate() 
    { 
        vector<vector<double> > res;
        for(auto &v : this->vertex)
            res.push_back(v.get_coordiante()); 
        return res; 
    };
    double con_Matrix(vector<Vertex<2> > &weight_points, vector<double> &quad_weights, int _i, int _j)
    {
        return 0 ;
    };
};
////////////////////////////////////////////////////////////////////////////////////////////////////////


class Q_Element;
class Q_Element : public Element<2>, public Square
{
public:
    Q_Element() :  Square() { this->index = -1 ;};
    /**
     * @brief 返回单元的四个顶点坐标
     * 
     * @return vector<vector<double> > 
     */
    vector<vector<double> > get_vertex_coordinate() 
    { 
        vector<vector<double> > res;
        for(auto &v : this->vertex)
            res.push_back(v.get_coordiante()); 
        return res; 
    } ;
    ///////////////////////////////////////////////////////////
    virtual double basic_function(int _i , double _xi, double _eta) = 0  ; //  第i个基函数在(_xi,_eta)点处的取值
    virtual vector<double> grad_basic_function(int _i, double _xi, double _eta) = 0; //第i个基函数在(_xi, _eta)处的梯度
    virtual vector<double> reference_grad_function(int _i, double _xi, double _eta)   //参考单元坐标转换函数的梯度
    {
        double partial_xi = 0 ;
        double partial_eta = 0 ;
        switch(_i)
        {
            case 0  :   partial_xi = (_eta -1) /4   ;   partial_eta = (_xi -1) /4 ;     break;
            case 1  :   partial_xi = -(_eta -1) /4  ;   partial_eta = -(_xi +1) /4;     break;
            case 2  :   partial_xi = (_eta +1) /4   ;   partial_eta = (_xi +1) /4;      break;
            case 3  :   partial_xi = -(_eta +1) /4  ;   partial_eta = -(_xi -1) /4;     break;
        };
        vector<double> res {partial_xi, partial_eta};
        return res;
    }
    /**double hx = (r_t.get_coordiante()[0] - l_b.get_coordiante()[0])/nx;
     * @brief 在点p处的jacobian矩阵
     * 
     * @param _p 
     * @return Eigen::Matrix2d 
     */
    Eigen::Matrix2d get_Jacobian(Vertex<2> _p)
    {
        Eigen::Matrix2d res = Eigen::Matrix2d::Zero(2,2);
        vector<double> point = _p.get_coordiante();
        double xi = point[0];
        double eta = point[1];
        const vector<vector<double> > vet = this->get_vertex_coordinate() ;
        for(int i =0; i < 4; i++)
        {
            res(0,0) += vet[i][0]* reference_grad_function(i,xi,eta)[0];
            res(0,1) += vet[i][1]* reference_grad_function(i,xi,eta)[0];
            res(1,0) += vet[i][0]* reference_grad_function(i,xi,eta)[1];
            res(1,1) += vet[i][1]* reference_grad_function(i,xi,eta)[1];
        }
        return res;
    }
    double con_Matrix(vector<Vertex<2> > &weight_points, vector<double> &quad_weights, int _i, int _j)
    {
        int n_wei_point = weight_points.size() ;
        double res = 0 ;
        for(int i =0; i < n_wei_point; i++)
        {
            Eigen::Matrix2d jacobian = this-> get_Jacobian(weight_points[i]);
            double det = jacobian(0,0)*jacobian(1,1) - jacobian(0,1)*jacobian(1,0);
            Eigen::Matrix2d inv_jacobian;
            inv_jacobian(0,0) = jacobian(1,1);
            inv_jacobian(0,1) = -jacobian(0,1);
            inv_jacobian(1,0) = -jacobian(1,0);
            inv_jacobian(1,1) = jacobian(1,1) ;
            inv_jacobian /= det;
            vector<double> grad_i = grad_basic_function(_i, weight_points[i].get_coordiante()[0],weight_points[i].get_coordiante()[1]);
            vector<double> grad_j = grad_basic_function(_j, weight_points[i].get_coordiante()[0],weight_points[i].get_coordiante()[1]);
            double par_fai_i_par_x = grad_i[0] * inv_jacobian(0,0) + grad_i[1] * inv_jacobian(0,1);
            double par_fai_j_par_x = grad_j[0] * inv_jacobian(0,0) + grad_j[1] * inv_jacobian(0,1);
            double par_fai_i_par_y= grad_i[0] * inv_jacobian(1,0) + grad_i[1] * inv_jacobian(1,1);
            double par_fai_j_par_y= grad_j[0] * inv_jacobian(1,0) + grad_j[1] * inv_jacobian(1,1);
            res += quad_weights[i] * det *(par_fai_i_par_x * par_fai_j_par_x + par_fai_i_par_y *par_fai_j_par_y);
        }
        return res ;
    }
};

class Q1_Element;
class Q1_Element : public Q_Element
{
public:
    Q1_Element() :  Q_Element() { this->index = -1 ; dofs.resize(4);};
    double basic_function(int _i , double _xi, double _eta)
    {
        double res = 0;
        switch(_i)
        {
            case 0  :   res = (_xi -1) *(_eta -1) /4;   break ;
            case 1  :   res = -(_xi +1) *(_eta -1) /4;  break;
            case 2  :   res = (_xi +1) *(_eta +1) /4;   break;
            case 3  :   res = -(_xi -1) *(_eta +1) /4;  break;
        }
        return res ;
    }
    vector<double> grad_basic_function(int _i, double _xi, double _eta)
    {
        return reference_grad_function(_i,_xi,_eta);
    }

};

template<int DIM, class element_type>
class Mesh;

template<int DIM, class element_type>
class Mesh
{
protected:
    vector<Dofs<DIM> > dofs;
    vector<element_type> elements;
    vector<size_t > boundary_dofs_index;
    vector<Vertex<DIM> > weight_point;
    vector<double> quad_weight;
public:
    void set_n_dofs( size_t _n) { dofs.resize(_n); };
    const vector<Dofs<DIM> >&  get_dofs() const{ return dofs; };
    void set_dofs(size_t _i, Dofs<DIM> _d) { dofs[_i] = _d; };
    void set_dofs(size_t _i, vector<double> _d, int _index) { dofs[_i].set_coordinate(_d) ;dofs[_i].set_index(_index); };
    Dofs<DIM> &  get_i_dofs(size_t _i) { return dofs[_i]; };
    void set_n_elements(size_t _n) { elements.resize(_n);};
    const vector<element_type >&  get_elements() const { return elements; };
    void set_elements(size_t _i, element_type _el) { elements[_i] = _el; };
    virtual element_type get_i_elements(size_t _i) = 0 ;
    void set_n_boundary_dofs(size_t _n) { boundary_dofs_index.resize(_n); };
    void set_boundary_dofs_index(vector<size_t> &vec) { boundary_dofs_index = vec ;};
    const vector<size_t> &  get_boundary_dofs_index() const {return boundary_dofs_index; };
    virtual void set_weight(int _acc) = 0 ;
    vector<Vertex<DIM> >& get_weight_point() { return weight_point;};
    vector<double>& get_quad_weight() { return quad_weight; };
};

template<class element_type>
class RectangleMesh : public Mesh<2, element_type>
{
protected:
    Vertex<2> left_bottom;
    Vertex<2> right_top;
    size_t nx;
    size_t ny;
public:
    RectangleMesh() = default ;
    RectangleMesh(vector<double> l_b, vector<double> r_t, size_t _nx, size_t _ny): nx(_nx),ny(_ny)
    {
        left_bottom.set_coordinate(l_b);
        right_top.set_coordinate(r_t);
    }
    Vertex<2>& get_left_bottom() { return left_bottom ;};
    void set_left_bottom(vector<double> l_b) { left_bottom.set_coordinate(l_b); };
    Vertex<2>& get_right_top() { return right_top ;};
    void set_right_top(vector<double> l_b) { right_top.set_coordinate(l_b); };
    size_t get_nx() { return nx;};
    void set_nx(size_t _nx) { nx = _nx; };
    size_t get_ny() { return ny;};
    void set_ny(size_t _ny) { ny = _ny;};
};

class Q1_RectangleMesh : public RectangleMesh<Q1_Element>
{
public:
    void set_weight(int _acc) ;
    Q1_RectangleMesh(vector<double> l_b, vector<double> r_t, size_t _nx, size_t _ny, int _acc) ;
    Q1_Element get_i_elements(size_t _i);
};

void Q1_RectangleMesh::set_weight(int _acc)
{
    string filename;
    string pathname ="GaussPntDoc/quadrilateral/";
    switch (_acc)
    {
        case 1: filename = "1.txt"; break;
        case 2:
        case 3: filename = "3.txt"; break;
        case 4:
        case 5: filename = "5.txt"; break;
        case 6:
        case 7: filename = "7.txt"; break;
        case 8:
        case 9: filename = "9.txt"; break;
        case 10:
        case 11: filename = "11.txt";break;    
        default: filename = "19.txt"; break;
    }
    filename = pathname + filename;
    ifstream is(filename.c_str());
    {
        int n_point;
        is >> n_point;
        weight_point.resize(n_point);
        quad_weight.resize(n_point);
        for(int i=0; i < n_point; i++)
        {
            vector<double> p(2);
            is >> p[0];
            is >> p[1];
            is >> quad_weight[i];
            weight_point[i].set_coordinate(p);
        }
    }
    is.close();
}

Q1_RectangleMesh::Q1_RectangleMesh(vector<double> l_b, vector<double> r_t, size_t _nx, size_t _ny, int _acc) : RectangleMesh(l_b, r_t, _nx, _ny)
{
    set_weight(_acc);
    /**
     * @brief 保存边界点编号
     * 
     */
    set_n_boundary_dofs(2*(_nx+_ny));
    size_t count = 0;
    for(size_t i =0; i < _nx+1; i++)
        boundary_dofs_index[count++] = i;
    for(size_t i =1; i < _ny ; i++)
    {
        boundary_dofs_index[count++] = (_nx+1)*i;
        boundary_dofs_index[count++] = (_nx+1)*(i+1) -1 ;
    }
    for(size_t i = 0; i < _nx +1; i++)
        boundary_dofs_index[count++] = (_nx +1)*_ny + i ;    
};
/**
 * @brief 获得第i个单元
 * 
 * @param _i 
 * @return Q1_Element 
 */
Q1_Element Q1_RectangleMesh:: get_i_elements(size_t _i)
{
    Q1_Element q_l;
    q_l.set_index(_i);
    size_t i = _i/nx ;
    size_t j = _i - i *nx ;
    vector<double> l_b = left_bottom.get_coordiante();
    vector<double> r_t = right_top.get_coordiante();
    double hx = (r_t[0] - l_b[0])/nx;
    double hy = (r_t[1] - l_b[1])/ny;
    vector<vector<double> > vertex;
    {
        Dofs<2> d;
        d.set_index(i*(nx+1)+j);
        vector<double> p = {j*hx+l_b[0],i*hy+l_b[1]};
        d.set_coordinate(p);
        q_l.set_dofs(d,0);
        vertex.push_back(p);
    }
    {
        Dofs<2> d;
        d.set_index(i*(nx+1)+j+1);
        vector<double> p = {j*hx+l_b[0]+hx,i*hy+l_b[1]};
        d.set_coordinate(p);
        q_l.set_dofs(d,1);
        vertex.push_back(p);
    }
    {
        Dofs<2> d;
        d.set_index((i+1)*(nx+1)+j+1);
        vector<double> p = {j*hx+l_b[0]+hx,i*hy+l_b[1]+hy};
        d.set_coordinate(p);
        q_l.set_dofs(d,2);
        vertex.push_back(p);
    }
    {
        Dofs<2> d;
        d.set_index((i+1)*(nx+1)+j);
        vector<double> p = {j*hx+l_b[0],i*hy+l_b[1]+hy};
        d.set_coordinate(p);
        q_l.set_dofs(d,3);
        vertex.push_back(p);
    }
    q_l.set_vertex(vertex);
    return q_l;
}

class Q1_Rectangel_fem_solution
{
private:
    Q1_RectangleMesh mesh;
    function<double(double,double)> f;
    function<double(double,double)> u;
public:
    Q1_Rectangel_fem_solution(vector<double> l_b, vector<double> r_t, size_t _nx, size_t _ny, int _acc, 
        function<double(double,double)> _f, function<double(double,double)> _u) : mesh(l_b,r_t,_nx,_ny,_acc), f(_f), u(_u) { };
    void solve();
};

void Q1_Rectangel_fem_solution::solve()
{
    size_t nx = this->mesh.get_nx();
    size_t ny = this->mesh.get_ny();
    Vertex<2> l_b = this->mesh.get_left_bottom();
    Vertex<2> r_t = this->mesh.get_right_top();
    vector<Vertex<2> > weight_point = this->mesh.get_weight_point(); 
    vector<double> quad_weight = this->mesh.get_quad_weight(); 
    const vector<size_t> boundary_index = this->mesh.get_boundary_dofs_index(); 
    size_t n_ele = nx * ny ;
    size_t n_dofs = (nx +1) * (ny +1);
    Eigen::MatrixXd A = Eigen::MatrixXd::Zero(n_dofs, n_dofs);
    Eigen::VectorXd F = Eigen::VectorXd::Zero(n_dofs);
    int weight_size = weight_point.size();
    double hx = (r_t.get_coordiante()[0] - l_b.get_coordiante()[0])/nx;
    double hy = (r_t.get_coordiante()[1] - l_b.get_coordiante()[1])/ny;
    for(size_t i =0; i < n_ele; i++)
    {
        Q1_Element Q1_e = this->mesh.get_i_elements(i);
        for(int j = 0; j < 4; j++)
            for(int k =0; k <4; k++)
            {
                A(Q1_e.get_i_dofs(j).get_index(), Q1_e.get_i_dofs(k).get_index() ) +=
                    Q1_e.con_Matrix(weight_point,quad_weight,j,k);
            }
        vector<vector<double> > vet = Q1_e.get_vertex_coordinate();
        double x, y ;
        for(int j =0; j < weight_size; j++)
        {
            for(int k =0 ; k < 4; k++)
            {
                Eigen::Matrix2d jacobi = Q1_e.get_Jacobian(weight_point[j]);
                double det = jacobi(0,0)*jacobi(1,1) - jacobi(1,0)*jacobi(0,1) ;
                x = (weight_point[j].get_coordiante()[0] + 1 ) / 2 * hx + vet[0][0] ;
                y = (weight_point[j].get_coordiante()[1] + 1 ) / 2 * hy + vet[0][1] ;
                F(Q1_e.get_i_dofs(k).get_index()) += quad_weight[j] * Q1_e.basic_function(k, weight_point[j].get_coordiante()[0],
                    weight_point[j].get_coordiante()[1])*det *(this->f(x,y));
            }                
        }
    }
    ////////////////////////////////////////////////////////////////
    //到此为止，把矩阵拼装完成。接下来，需要对边界进行处理
    
    for(auto &index : boundary_index)
    {
        //先获得该点的坐标
        size_t ii = index/(nx+1);
        size_t jj = index - ii *(nx+1);
        double xi = ii * hx + l_b.get_coordiante()[0];
        double yi = jj * hy + l_b.get_coordiante()[1];
        double value = this->u(xi,yi);
        for(size_t i = 0; i < n_dofs; i++)
        {
            F(i) -= A(i,index)*value;
            A(i,index) = 0 ;
        }
    }
    for(auto & index : boundary_index)
    {
        size_t ii = index/(nx+1);
        size_t jj = index - ii *(nx+1);
        double xi = ii * hx + l_b.get_coordiante()[0];
        double yi = jj * hy + l_b.get_coordiante()[1];
        double value = this->u(xi,yi);
        for(size_t i = 0; i < index; i++)
        {
            A(index,i) = 0 ;
        }
        for(size_t i = index+1; i < n_dofs; i++)
        {
            A(index,i) = 0 ;
        }
        F(index) = value*A(index,index);
    }
    Eigen::VectorXd X;
    X = A.colPivHouseholderQr().solve(F);
    ofstream os("res_data.txt");
    os << "nx:   " << nx <<"ny:    " << ny << std::endl;
    for(size_t j = 0; j < ny+1; j++)
        for(size_t i =0; i < nx+1; i++)
        {
            os <<  setiosflags(std::ios::right);
            os << setw(15) << l_b.get_coordiante()[0]+ i * hx << setw(15) << 
                l_b.get_coordiante()[1]+ j*hy <<setw(15) << X(j*(nx+1)+i) << std::endl ;
        }
    os.close();
    std::cout << X << std::endl;
}   

#endif